import os
import sys

current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
ch03_dir = os.path.join(parent_dir, "ch03")
sys.path.append(ch03_dir)

import torch
import torch.nn as nn
import torch.nn.functional as F

from attention import MultiHeadAttention


class GPTModel(nn.Module):
    """GPT模型实现，基于Transformer架构。

    Attributes:
        tok_emb (nn.Embedding): 词嵌入层
        pos_emb (nn.Embedding): 位置嵌入层
        drop_emb (nn.Dropout): 嵌入层后的dropout
        trf_blocks (nn.Sequential): Transformer块序列
        final_norm (LayerNorm): 最终层归一化
        out_head (nn.Linear): 输出层
    """

    def __init__(self, config: dict) -> None:
        """初始化GPT模型。

        Args:
            config (dict): 模型配置字典，包含以下键：
                - vocab_size: int, 词汇表大小
                - emb_dim: int, 嵌入维度
                - context_length: int, 上下文长度
                - drop_rate: float, dropout比率
                - n_layers: int, Transformer层数
                - n_heads: int, 注意力头数
                - qkv_bias: bool, 是否使用QKV偏置
        """
        super().__init__()
        self.tok_emb = nn.Embedding(config["vocab_size"], config["emb_dim"])
        self.pos_emb = nn.Embedding(config["context_length"], config["emb_dim"])
        self.drop_emb = nn.Dropout(config["drop_rate"])

        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(config) for _ in range(config["n_layers"])]
        )

        self.final_norm = LayerNorm(config["emb_dim"])
        self.out_head = nn.Linear(config["emb_dim"], config["vocab_size"], bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 [batch_size, seq_len]

        Returns:
            torch.Tensor: 输出logits，形状为 [batch_size, seq_len, vocab_size]
        """
        batch_size, seq_len = x.shape
        tok_embs = self.tok_emb(x)
        pos_embs = self.pos_emb(torch.arange(seq_len, device=x.device))
        x = tok_embs + pos_embs
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

class TransformerBlock(nn.Module):
    """Transformer块，包含多头注意力、前馈网络和层归一化。

    Attributes:
        attn (MultiHeadAttention): 多头注意力层
        ln1 (LayerNorm): 第一个层归一化
        ff (FeedForward): 前馈网络
        ln2 (LayerNorm): 第二个层归一化
        dropout (nn.Dropout): dropout层
    """

    def __init__(self, config: dict) -> None:
        """初始化Transformer块。

        Args:
            config (dict): 配置字典，包含以下键：
                - emb_dim: int, 嵌入维度
                - context_length: int, 上下文长度
                - n_heads: int, 注意力头数
                - drop_rate: float, dropout比率
                - qkv_bias: bool, 是否使用QKV偏置
        """
        super().__init__()
        self.attn = MultiHeadAttention(
            hidden_dim=config["emb_dim"],
            d_out=config["emb_dim"],
            context_length=config["context_length"],
            num_heads=config["n_heads"],
            dropout=config["drop_rate"],
            qkv_bias=config["qkv_bias"]
        )
        self.ln1 = LayerNorm(config["emb_dim"])
        self.ff = FeedForward(config)
        self.ln2 = LayerNorm(config["emb_dim"])
        self.dropout = nn.Dropout(config["drop_rate"])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 [batch_size, seq_len, emb_dim]

        Returns:
            torch.Tensor: 输出张量，形状与输入相同
        """
        shortcut = x
        x = self.ln1(x)
        x = self.attn(x)
        x = self.dropout(x)
        x = x + shortcut

        shortcut = x
        x = self.ln2(x)
        x = self.ff(x)
        x = self.dropout(x)
        x = x + shortcut
        return x

class LayerNorm(nn.Module):
    """层归一化实现。

    Attributes:
        eps (float): 数值稳定性参数
        scale (nn.Parameter): 缩放参数
        shift (nn.Parameter): 平移参数
    """

    def __init__(self, emb_dim: int, eps: float = 1e-5) -> None:
        """初始化层归一化。

        Args:
            emb_dim (int): 嵌入维度
            eps (float, optional): 数值稳定性参数，默认为1e-5
        """
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 [..., emb_dim]

        Returns:
            torch.Tensor: 归一化后的张量，形状与输入相同
        """
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)
        return self.scale * (x - mean) / (std + self.eps) + self.shift

class GELU(nn.Module):
    """GELU激活函数实现。

    GELU (Gaussian Error Linear Unit) 是一种常用的激活函数，在Transformer架构中广泛使用。
    """

    def __init__(self) -> None:
        """初始化GELU激活函数。"""
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。

        Args:
            x (torch.Tensor): 输入张量

        Returns:
            torch.Tensor: 应用GELU激活后的张量
        """
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) * 
            (x + 0.044715 * torch.pow(x, 3))
        ))
    
class FeedForward(nn.Module):
    """前馈网络实现。

    Attributes:
        layers (nn.Sequential): 包含线性层和激活函数的序列
    """

    def __init__(self, cfg: dict) -> None:
        """初始化前馈网络。

        Args:
            cfg (dict): 配置字典，需要包含emb_dim键
        """
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"])
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播。

        Args:
            x (torch.Tensor): 输入张量，形状为 [..., emb_dim]

        Returns:
            torch.Tensor: 输出张量，形状与输入相同
        """
        return self.layers(x)


def generate_text_simple(
    model: GPTModel,
    idx: torch.Tensor,
    max_new_tokens: int,
    context_length: int
) -> torch.Tensor:
    """使用模型生成文本的简单实现。

    Args:
        model (GPTModel): 训练好的GPT模型
        idx (torch.Tensor): 初始输入序列，形状为 [batch_size, seq_len]
        max_new_tokens (int): 要生成的新token数量
        context_length (int): 上下文窗口长度

    Returns:
        torch.Tensor: 生成的完整序列，形状为 [batch_size, seq_len + max_new_tokens]
    """
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_length:]
        with torch.no_grad():
            logits = model(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.argmax(probs, dim=-1, keepdim=True)
            idx = torch.cat((idx, idx_next), dim=1)
    return idx